from ultralytics import YOLO
from collections import Counter
import numpy as np


class YOLOProcessor:
    def __init__(self, model_path):
        self.model = YOLO(model_path)

    def process_image(self, image_path):
        """处理图片"""
        return self.model(image_path)

    def process_frame(self, frame):
        """处理视频帧"""
        return self.model(frame)

    def get_class_names(self):
        """获取类别名称"""
        return self.model.names

    def extract_detections(self, results):
        """从结果中提取检测信息"""
        detections = []
        if results[0].boxes is not None:
            boxes = results[0].boxes
            for i in range(len(boxes)):
                cls = int(boxes.cls[i].item())
                conf = boxes.conf[i].item()
                detections.append({
                    'class': cls,
                    'confidence': conf,
                    'class_name': self.model.names[cls]
                })
        return detections

    def format_detection_results(self, results):
        """格式化检测结果"""
        if results.boxes is None or len(results.boxes) == 0:
            return '未检测到目标', '未检测到任何目标物体'

        boxes = results.boxes
        detections = []

        # 收集所有检测结果
        for i in range(len(boxes)):
            cls = int(boxes.cls[i].item())
            conf = boxes.conf[i].item()
            class_name = self.model.names[cls]
            detections.append({
                'class_id': cls,
                'class_name': class_name,
                'confidence': conf
            })

        # 统计信息
        total_objects = len(detections)
        class_counter = Counter([d['class_name'] for d in detections])
        avg_confidence = np.mean([d['confidence'] for d in detections])

        # 更新统计标签
        stats_text = f"""
总检测数量: {total_objects}
平均置信度: {avg_confidence:.3f}
检测类别数: {len(class_counter)}
        """

        # 更新详细结果
        results_text = f"检测结果统计:\n"
        results_text += f"总目标数量: {total_objects}\n"
        results_text += f"平均置信度: {avg_confidence:.3f}\n\n"

        results_text += "各类别数量:\n"
        for class_name, count in class_counter.items():
            results_text += f"  {class_name}: {count}个\n"

        results_text += f"\n详细检测结果:\n"
        results_text += "-" * 40 + "\n"

        for i, detection in enumerate(detections, 1):
            results_text += (f"目标 {i}: {detection['class_name']} "
                             f"(置信度: {detection['confidence']:.3f})\n")

        return stats_text, results_text

    def format_video_results(self, detections):
        """格式化视频检测结果"""
        if not detections:
            stats_text = "实时检测: 未检测到目标"
            results_text = "当前帧未检测到任何目标"
        else:
            # 统计信息
            total_objects = len(detections)
            class_counter = Counter([d['class_name'] for d in detections])
            avg_confidence = np.mean([d['confidence'] for d in detections])

            stats_text = f"""
实时检测:
总数量: {total_objects}
平均置信度: {avg_confidence:.3f}
类别数: {len(class_counter)}
            """

            # 详细结果
            results_text = "实时检测结果:\n\n"
            results_text += "各类别数量:\n"
            for class_name, count in class_counter.most_common():
                results_text += f"  {class_name}: {count}个\n"

            results_text += f"\n最近检测:\n"
            for detection in detections[:5]:  # 显示最近5个检测
                results_text += (f"  {detection['class_name']} "
                                 f"({detection['confidence']:.3f})\n")

        return stats_text, results_text
